import random
from enum import IntEnum

import pygame
import pymunk
import pymunk.pygame_util

difficulty_factor = 2.0  # 难度因子


# 水果枚举，数字代表水果级别
class FruitType(IntEnum):
    CHERRY = 0  # 樱桃(最小)
    STRAWBERRY = 1  # 草莓
    GRAPE = 2  # 葡萄
    ORANGE = 3  # 橙子
    PERSIMMON = 4  # 柿子
    APPLE = 5  # 苹果
    PEAR = 6  # 梨
    PEACH = 7  # 桃子
    PINEAPPLE = 8  # 菠萝
    MELON = 9  # 哈密瓜
    WATERMELON = 10  # 西瓜(最大)


# 水果属性: (半径, 质量, 颜色)
# 水果属性: (半径, 质量, 颜色)
FRUIT_PROPERTIES = {
    FruitType.CHERRY: (10 * difficulty_factor, 1, (255, 0, 0)),  # 红色 (原15)
    FruitType.STRAWBERRY: (13 * difficulty_factor, 2, (255, 0, 127)),  # 粉红色 (原20)
    FruitType.GRAPE: (16 * difficulty_factor, 4, (128, 0, 128)),  # 紫色 (原25)
    FruitType.ORANGE: (20 * difficulty_factor, 8, (255, 165, 0)),  # 橙色 (原30)
    FruitType.PERSIMMON: (25 * difficulty_factor, 16, (255, 97, 0)),  # 深橙色 (原35)
    FruitType.APPLE: (31 * difficulty_factor, 32, (255, 0, 0)),  # 红色 (原40)
    FruitType.PEAR: (40 * difficulty_factor, 64, (173, 255, 47)),  # 黄绿色 (原45)
    FruitType.PEACH: (50 * difficulty_factor, 128, (255, 218, 185)),  # 桃色 (原50)
    FruitType.PINEAPPLE: (64 * difficulty_factor, 256, (255, 255, 0)),  # 黄色 (原55)
    FruitType.MELON: (80 * difficulty_factor, 512, (152, 251, 152)),  # 淡绿色 (原60)
    FruitType.WATERMELON: (100 * difficulty_factor, 1024, (0, 128, 0)),  # 深绿色 (原70)
}


class WatermelonGame:
    def __init__(self, width=400, height=600):
        # 初始化pygame和pymunk
        pygame.init()
        self.width = width
        self.height = height
        self.screen = pygame.Surface((width, height))

        # 创建物理空间
        self.space = pymunk.Space()
        self.space.gravity = (0, 500)  # 向下的重力
        self.draw_options = pymunk.pygame_util.DrawOptions(self.screen)

        # 游戏参数
        self.next_fruit = None
        self.current_fruit = None
        self.waiting_for_drop = True
        self.game_over = False
        self.victory = False  # 新增：胜利状态标志
        self.score = 0

        # 创建边界墙
        self._create_walls()

        # 当前掉落水果的位置
        self.drop_position = width // 2

        # 准备下一个水果
        self.prepare_next_fruit()

        # 碰撞处理
        self.collision_handler = self.space.add_collision_handler(0, 0)
        self.collision_handler.post_solve = self.handle_collision

        # 要合并的水果对
        self.fruits_to_merge = set()

        # 宽限期计时器
        self.grace_period = 0

        # 添加新水果冷却时间
        self.refresh_cooldown = 0
        self.refresh_time = 30  # 等待30帧刷新新水果

    def _create_walls(self):
        """创建游戏边界"""
        walls = [
            # 底部
            [(0, self.height), (self.width, self.height)],
            # 左侧
            [(0, 0), (0, self.height)],
            # 右侧
            [(self.width, 0), (self.width, self.height)],
        ]

        for wall in walls:
            body = pymunk.Body(body_type=pymunk.Body.STATIC)
            shape = pymunk.Segment(body, wall[0], wall[1], 5)
            shape.elasticity = 0.4
            shape.friction = 0.5
            self.space.add(body, shape)

    def prepare_next_fruit(self):
        """准备下一个水果"""
        # 随机选择前3种最小的水果
        fruit_type = random.choice(list(FruitType)[:3])
        self.next_fruit = fruit_type
        if self.current_fruit is None:
            self.current_fruit = self.next_fruit
            self.next_fruit = random.choice(list(FruitType)[:3])

    def create_fruit(self, fruit_type, position):
        """创建一个新水果"""
        radius, mass, color = FRUIT_PROPERTIES[fruit_type]
        moment = pymunk.moment_for_circle(mass, 0, radius)
        body = pymunk.Body(mass, moment)
        body.position = position

        shape = pymunk.Circle(body, radius)
        shape.elasticity = 0.4
        shape.friction = 0.5
        shape.collision_type = 0

        # 存储水果类型和颜色信息
        shape.fruit_type = fruit_type
        shape.color = color

        self.space.add(body, shape)
        return shape

    def handle_collision(self, arbiter, space, data):
        """碰撞后的回调处理"""
        shapes = arbiter.shapes
        if len(shapes) != 2:
            return True

        s1, s2 = shapes

        # 检查两个水果是否相同类型
        if (
            hasattr(s1, "fruit_type")
            and hasattr(s2, "fruit_type")
            and s1.fruit_type == s2.fruit_type
            and s1.fruit_type != FruitType.WATERMELON
        ):  # 西瓜不能再合并
            # 将这对水果添加到等待合并的集合中
            self.fruits_to_merge.add((s1, s2))

        return True

    def process_merges(self):
        """处理所有待合并的水果对"""
        if not self.fruits_to_merge:
            return

        processed = set()
        for s1, s2 in self.fruits_to_merge:
            # 如果已经处理过这个水果，则跳过
            if s1 in processed or s2 in processed or not s1.body or not s2.body:
                continue

            # 计算新水果的位置(两个水果的中点)
            pos = (
                (s1.body.position.x + s2.body.position.x) / 2,
                (s1.body.position.y + s2.body.position.y) / 2,
            )

            # 移除旧水果
            self.space.remove(s1, s1.body)
            self.space.remove(s2, s2.body)
            processed.add(s1)
            processed.add(s2)

            # 创建新的更大的水果
            new_type = FruitType(s1.fruit_type + 1)
            self.create_fruit(new_type, pos)

            # 检查是否合成了西瓜
            if new_type == FruitType.WATERMELON:
                self.victory = True  # 设置胜利标志

            # 更新分数 - 使用指数奖励
            self.score += 2 ** (new_type.value)  # 2的幂次奖励

        # 清空合并列表
        self.fruits_to_merge.clear()

    def update_drop_position(self, dx):
        """更新水果的水平投放位置"""
        self.drop_position += dx
        # 确保不超出边界
        radius, _, _ = FRUIT_PROPERTIES[self.current_fruit]
        self.drop_position = max(
            radius + 10, min(self.width - radius - 10, self.drop_position)
        )

    def drop_fruit(self):
        """投放当前水果"""
        if not self.waiting_for_drop:
            return

        # 计算场上水果数量并施加惩罚（仅在投放时生效）
        fruit_count = sum(
            1 for body in self.space.bodies if body.body_type == pymunk.Body.DYNAMIC
        )
        # 惩罚系数可以根据需要调整
        fruit_penalty = (fruit_count**2) * 0.02  # 可以适当增大系数，因为频率变低了
        self.score -= fruit_penalty

        # 获取当前水果的半径
        radius, _, _ = FRUIT_PROPERTIES[self.current_fruit]

        # 投放位置在顶部
        position = (self.drop_position, radius + 10)
        self.create_fruit(self.current_fruit, position)

        # 更新状态
        self.waiting_for_drop = False
        self.current_fruit = self.next_fruit
        self.prepare_next_fruit()

        # 设置宽限期为60帧(1秒)
        self.grace_period = 60

        # 设置新水果冷却时间
        self.refresh_cooldown = self.refresh_time

    def check_game_over(self):
        """检查游戏是否结束"""
        # 设置绝对高度限制 - 这个限制不受宽限期影响
        absolute_threshold_y = 30  # 绝对高度限制，比投放位置更高

        # 首先检查是否有水果超过绝对高度限制
        for body in self.space.bodies:
            if body.body_type == pymunk.Body.DYNAMIC:
                # 检查是否超过绝对高度限制
                if body.position.y < absolute_threshold_y:
                    return True  # 立即结束游戏，不考虑宽限期

        # 宽限期内不进行常规结束检查
        if self.grace_period > 0:
            return False

        # 常规结束检查：如果有水果在顶部区域停留过长时间
        threshold_y = 120  # 常规顶部区域高度

        for body in self.space.bodies:
            if body.body_type == pymunk.Body.DYNAMIC:
                # 只检查动态物体
                if body.position.y < threshold_y and abs(body.velocity.y) < 20:
                    # 如果水果速度很小且在顶部区域
                    return True

        # 检查水果总数，防止堆积过多导致性能问题
        fruit_count = sum(
            1 for body in self.space.bodies if body.body_type == pymunk.Body.DYNAMIC
        )
        if fruit_count > 30:  # 如果场上水果过多，结束游戏
            return True

        return False

    def step(self, action=None):
        """游戏前进一步"""
        if action is not None:
            # 简化为只处理投放动作，移除左移右移
            if action == 2:  # 投放
                self.drop_fruit()

        # 更新物理模拟
        self.space.step(1 / 60.0)

        # 更新宽限期计时器
        if self.grace_period > 0:
            self.grace_period -= 1

        # 更新新水果冷却时间
        if not self.waiting_for_drop and self.refresh_cooldown > 0:
            self.refresh_cooldown -= 1
            if self.refresh_cooldown <= 0:
                self.waiting_for_drop = True

        # 处理水果合并
        self.process_merges()

        # 移除此处的惩罚代码

        # 检查游戏是否结束
        self.game_over = self.check_game_over() or self.victory

        return self.game_over, self.victory  # 返回游戏结束状态和胜利状态

    def render(self):
        """渲染游戏画面"""
        # 清空屏幕
        self.screen.fill((255, 255, 255))

        # 绘制所有物体
        for body in self.space.bodies:
            if body.body_type != pymunk.Body.STATIC:  # 不渲染静态物体(墙)
                for shape in body.shapes:
                    if hasattr(shape, "color"):
                        # 绘制水果
                        pygame.draw.circle(
                            self.screen,
                            shape.color,
                            (int(body.position.x), int(body.position.y)),
                            int(shape.radius),
                        )

        # 当前水果 - 只在可投放状态显示
        if self.waiting_for_drop:
            current_radius, _, current_color = FRUIT_PROPERTIES[self.current_fruit]
            # 在投放位置显示当前水果
            pygame.draw.circle(
                self.screen,
                current_color,
                (self.drop_position, current_radius + 10),
                current_radius,
            )

        # 始终在右上角显示下一个水果预览
        next_radius, _, next_color = FRUIT_PROPERTIES[self.next_fruit]
        preview_radius = min(next_radius * 0.7, 30)  # 预览尺寸略小
        pygame.draw.circle(
            self.screen, next_color, (self.width - 50, 50), int(preview_radius)
        )

        # 在右上角添加"Next:"标签
        font = pygame.font.SysFont(None, 24)
        text = font.render("Next", True, (0, 0, 0))
        self.screen.blit(text, (self.width - 80, 30))

        # 返回游戏画面
        return self.screen.copy()

    def reset(self):
        """重置游戏状态"""
        # 清空物理空间中所有动态物体
        for body in list(self.space.bodies):
            if body.body_type == pymunk.Body.DYNAMIC:
                self.space.remove(body, *body.shapes)

        # 重置游戏参数
        self.waiting_for_drop = True
        self.game_over = False
        self.score = 0
        self.fruits_to_merge.clear()
        self.drop_position = self.width // 2

        # 添加这一行：设置初始宽限期
        self.grace_period = 60

        # 重置胜利状态
        self.victory = False

        # 准备新水果
        self.prepare_next_fruit()

        return self.render()
